{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Supervised relation extraction with LukeForEntityPairClassification.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "6c739827bd3c48e68775a3951bebe7c9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_46f830096d7449c680d69b6111e73dfe",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_870102d2c71b473e9f57f62de01636d3",
              "IPY_MODEL_2759d4db7a2b435dba9d6c6b923863b2"
            ]
          }
        },
        "46f830096d7449c680d69b6111e73dfe": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "870102d2c71b473e9f57f62de01636d3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_08e58a87449945819d0d3fe3a64636a5",
            "_dom_classes": [],
            "description": "Validation sanity check:   0%",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 2,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_78993027eda248ccb5753ae2629e8b63"
          }
        },
        "2759d4db7a2b435dba9d6c6b923863b2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_206e3619d1d8476bbef8a5aa9d1c7b5c",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/2 [19:15&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_8aedf1e4705a4456ba9dbfb9acefae6a"
          }
        },
        "08e58a87449945819d0d3fe3a64636a5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "78993027eda248ccb5753ae2629e8b63": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "206e3619d1d8476bbef8a5aa9d1c7b5c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "8aedf1e4705a4456ba9dbfb9acefae6a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "60e547a3f89b4ff6b21d22f1e4a101f8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_c32a8d9b12704d42a2d784c59a2c666f",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_57371a01d21d474887f025f4a7579acf",
              "IPY_MODEL_f4528f4bd08a41f7966768850ef2e6d0"
            ]
          }
        },
        "c32a8d9b12704d42a2d784c59a2c666f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "57371a01d21d474887f025f4a7579acf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_fcf36c3be3044f38a2ee84c5c204e303",
            "_dom_classes": [],
            "description": "Epoch 4:   5%",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 2888,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 140,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_07ce5de710e446649fbd5dc84d393cb0"
          }
        },
        "f4528f4bd08a41f7966768850ef2e6d0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9e8ed2c19561424c972bac52133019bf",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 140/2888 [01:19&lt;26:08,  1.75it/s, loss=0.0394, v_num=1nbe]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_1069cb96cfc7413c86e49569d38e946e"
          }
        },
        "fcf36c3be3044f38a2ee84c5c204e303": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "07ce5de710e446649fbd5dc84d393cb0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "9e8ed2c19561424c972bac52133019bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "1069cb96cfc7413c86e49569d38e946e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d5d713d04c26461fb098dd15145084c2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_69310f37f6f8416cb93426339816405b",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_80d1309b5e0e4c8e8433402ec3faeb78",
              "IPY_MODEL_62b8e61826804383a726c459852d47f9"
            ]
          }
        },
        "69310f37f6f8416cb93426339816405b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "80d1309b5e0e4c8e8433402ec3faeb78": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_390bf45600c2402abaf8baa1459ee4f9",
            "_dom_classes": [],
            "description": "Validating: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 963,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 963,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_5db797f8ba184f718a11b72b5e798319"
          }
        },
        "62b8e61826804383a726c459852d47f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_3d09aea45dfc43f6ac7f87ed29bf6918",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 963/963 [01:27&lt;00:00, 11.57it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_152564b90b62466aab21e6f7744d7024"
          }
        },
        "390bf45600c2402abaf8baa1459ee4f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "5db797f8ba184f718a11b72b5e798319": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3d09aea45dfc43f6ac7f87ed29bf6918": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "152564b90b62466aab21e6f7744d7024": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a22625735f704a8f9751e02b11edd52a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_7f3319a4c4bc45479342a111e47d7afb",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_e879d2939d1048acb218ca36661ae997",
              "IPY_MODEL_6bcc6aa979ce4842a41aa7733a4766c8"
            ]
          }
        },
        "7f3319a4c4bc45479342a111e47d7afb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "e879d2939d1048acb218ca36661ae997": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_c4cf2946d1e3425b8c041e901f9580d9",
            "_dom_classes": [],
            "description": "Validating: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 963,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 963,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_464a61cb2b4b4eccbeb61fdea07de12c"
          }
        },
        "6bcc6aa979ce4842a41aa7733a4766c8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_ad9ba418e52248a09313c8f13ac7bfe1",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 963/963 [01:27&lt;00:00, 11.53it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d7eac30805ce40bb95e4c679ba4126af"
          }
        },
        "c4cf2946d1e3425b8c041e901f9580d9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "464a61cb2b4b4eccbeb61fdea07de12c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ad9ba418e52248a09313c8f13ac7bfe1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "d7eac30805ce40bb95e4c679ba4126af": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "466d5d0c87a0493c9a77e4c01233d94e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_79d771361fb84c62a18c93f9e70946cf",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_d8d4525a722a4ba686037016eb8524da",
              "IPY_MODEL_35e227d4e3af46c5856c32587ff18860"
            ]
          }
        },
        "79d771361fb84c62a18c93f9e70946cf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "d8d4525a722a4ba686037016eb8524da": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_12d42fdf1c6e4050ad48907759d3dbea",
            "_dom_classes": [],
            "description": "Validating: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 963,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 963,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_291b1496ede44d738f63f94c976f6f25"
          }
        },
        "35e227d4e3af46c5856c32587ff18860": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_4ce1f43afd33491fa22dd77403cf08ba",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 963/963 [01:26&lt;00:00, 11.49it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_0b5233fe3cc047339d33cb2b91ec0908"
          }
        },
        "12d42fdf1c6e4050ad48907759d3dbea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "291b1496ede44d738f63f94c976f6f25": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4ce1f43afd33491fa22dd77403cf08ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "0b5233fe3cc047339d33cb2b91ec0908": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4efd9113c15d443aa5fadb039e0a31d3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_ddadd8532e3a4ca9925d748144f020bd",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_3883b380d960454ab297f64914d8af25",
              "IPY_MODEL_c3a8c773516c46139dae4dd17e932528"
            ]
          }
        },
        "ddadd8532e3a4ca9925d748144f020bd": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "3883b380d960454ab297f64914d8af25": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_6f83ac1276254e939157238317531e24",
            "_dom_classes": [],
            "description": "Validating: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 963,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 963,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_b8d8686696334b1089ae1356c17e51c4"
          }
        },
        "c3a8c773516c46139dae4dd17e932528": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_eb8531165ba246fe9afb40c498de7d17",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 963/963 [01:27&lt;00:00, 11.56it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_b1c31fbb42ea48f3bbe19820f868991e"
          }
        },
        "6f83ac1276254e939157238317531e24": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "b8d8686696334b1089ae1356c17e51c4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "eb8531165ba246fe9afb40c498de7d17": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "b1c31fbb42ea48f3bbe19820f868991e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "61bca5692a0146ebad95a0b43cc50eaa": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_23a961a7ebc241b393e408b04d9e5d1d",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_ff3c708e85b14e5dbbd3cbebc8f1c6cd",
              "IPY_MODEL_e17ac0baf40943c5b0ba3c329c22e992"
            ]
          }
        },
        "23a961a7ebc241b393e408b04d9e5d1d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ff3c708e85b14e5dbbd3cbebc8f1c6cd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_68fb7646daf141ad9eb00159d7ae8f40",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1204,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1204,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_9e9c310829134f8b9825c1a899e88746"
          }
        },
        "e17ac0baf40943c5b0ba3c329c22e992": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_ce141efe165d4aa0817c0e7e518039df",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1204/1204 [12:14&lt;00:00,  1.64it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_42f47d64f2924e7693e16d11a5457136"
          }
        },
        "68fb7646daf141ad9eb00159d7ae8f40": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "9e9c310829134f8b9825c1a899e88746": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ce141efe165d4aa0817c0e7e518039df": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "42f47d64f2924e7693e16d11a5457136": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/LUKE/Supervised_relation_extraction_with_LukeForEntityPairClassification.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EhOUh0qzFFje"
      },
      "source": [
        "In this notebook, we are going to fine-tune [`LukeForEntityPairClassification`](https://huggingface.co/transformers/model_doc/luke.html#lukeforentitypairclassification) on a supervised **relation extraction** dataset.\n",
        "\n",
        "The goal for the model is to predict, given a sentence and the character spans of two entities within the sentence, the relationship between the entities.\n",
        "\n",
        "The author of LUKE has fine-tuned this model on the [TACRED](https://nlp.stanford.edu/projects/tacred/) dataset, an important supervised relation extraction dataset by Stanford University, and obtains state-of-the-art results with it. \n",
        "\n",
        "* Paper: https://arxiv.org/abs/2010.01057\n",
        "* Original repository: https://github.com/studio-ousia/luke"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9sFqwmObswh9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5daef8ff-9d85-4264-92da-fe1abfa4ce6a"
      },
      "source": [
        "!pip install -q transformers "
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[K     |████████████████████████████████| 2.3MB 2.9MB/s \n",
            "\u001b[K     |████████████████████████████████| 901kB 17.7MB/s \n",
            "\u001b[K     |████████████████████████████████| 3.3MB 24.2MB/s \n",
            "\u001b[?25h"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9wWQBN-FSIN",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "acad943a-f4f6-4bc5-eaeb-f96598df2325"
      },
      "source": [
        "!pip install -q pytorch-lightning wandb"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[K     |████████████████████████████████| 808kB 2.9MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.8MB 16.1MB/s \n",
            "\u001b[K     |████████████████████████████████| 645kB 20.2MB/s \n",
            "\u001b[K     |████████████████████████████████| 829kB 19.6MB/s \n",
            "\u001b[K     |████████████████████████████████| 112kB 29.4MB/s \n",
            "\u001b[K     |████████████████████████████████| 10.6MB 24.0MB/s \n",
            "\u001b[K     |████████████████████████████████| 276kB 37.3MB/s \n",
            "\u001b[K     |████████████████████████████████| 102kB 10.0MB/s \n",
            "\u001b[K     |████████████████████████████████| 174kB 38.5MB/s \n",
            "\u001b[K     |████████████████████████████████| 133kB 40.7MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.3MB 35.2MB/s \n",
            "\u001b[K     |████████████████████████████████| 71kB 8.9MB/s \n",
            "\u001b[K     |████████████████████████████████| 143kB 30.2MB/s \n",
            "\u001b[K     |████████████████████████████████| 296kB 35.4MB/s \n",
            "\u001b[?25h  Building wheel for future (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for subprocess32 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement tensorboard~=2.5, but you'll have tensorboard 2.4.1 which is incompatible.\u001b[0m\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bUE24XG63vPM"
      },
      "source": [
        "## Read in data\n",
        "\n",
        "Let's download the data from the web, hosted on Dropbox."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XeKagbUCNj05"
      },
      "source": [
        "import requests, zipfile, io\n",
        "\n",
        "def download_data():\n",
        "    url = \"https://www.dropbox.com/s/izi2x4sjohpzoot/relation_extraction_dataset.zip?dl=1\"\n",
        "    r = requests.get(url)\n",
        "    z = zipfile.ZipFile(io.BytesIO(r.content))\n",
        "    z.extractall()\n",
        "\n",
        "download_data()"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TG7bi0fYOfw2"
      },
      "source": [
        "Each row in the dataframe consists of a news article, and a sentence in which a certain relationship was found (just as \"invested_in\", or \"founded_by\"). There were some patterns used to gather the data, so it might contain some noise. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 287
        },
        "id": "hDkptorP9Koh",
        "outputId": "e46ddab9-470a-45fa-877a-a2f5abadf835"
      },
      "source": [
        "import pandas as pd\n",
        "\n",
        "df = pd.read_pickle(\"relation_extraction_dataset.pkl\")\n",
        "df.reset_index(drop=True, inplace=True)\n",
        "df.head()"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>end_idx</th>\n",
              "      <th>entities</th>\n",
              "      <th>entity_spans</th>\n",
              "      <th>match</th>\n",
              "      <th>original_article</th>\n",
              "      <th>sentence</th>\n",
              "      <th>start_idx</th>\n",
              "      <th>string_id</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1024</td>\n",
              "      <td>[Lilium, Baillie Gifford]</td>\n",
              "      <td>[[3, 9], [151, 166]]</td>\n",
              "      <td>raising $35</td>\n",
              "      <td>Happy Friday!\\n\\nWe sincerely hope you and you...</td>\n",
              "      <td>3) Lilium, a German startup that’s making an a...</td>\n",
              "      <td>1013</td>\n",
              "      <td>invested_in</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1762</td>\n",
              "      <td>[Facebook ’s, Giphy]</td>\n",
              "      <td>[[92, 102], [148, 153]]</td>\n",
              "      <td>acquisition</td>\n",
              "      <td>Happy Friday!\\n\\nWe sincerely hope you and you...</td>\n",
              "      <td>Meanwhile, the UK’s watchdog on Friday announc...</td>\n",
              "      <td>1751</td>\n",
              "      <td>acquired_by</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>2784</td>\n",
              "      <td>[Global-e, Vitruvian Partners]</td>\n",
              "      <td>[[27, 35], [94, 112]]</td>\n",
              "      <td>raised $60</td>\n",
              "      <td>Happy Friday!\\n\\nWe sincerely hope you and you...</td>\n",
              "      <td>Israeli e-commerce startup Global-e has raised...</td>\n",
              "      <td>2774</td>\n",
              "      <td>invested_in</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>680</td>\n",
              "      <td>[Joris Van Der Gucht, Silverfin]</td>\n",
              "      <td>[[0, 19], [35, 44]]</td>\n",
              "      <td>founder</td>\n",
              "      <td>Hg is a leading investor in tax and accounting...</td>\n",
              "      <td>Joris Van Der Gucht, co-founder at Silverfin c...</td>\n",
              "      <td>673</td>\n",
              "      <td>founded_by</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>2070</td>\n",
              "      <td>[Tim Vandecasteele, Silverfin]</td>\n",
              "      <td>[[0, 17], [71, 80]]</td>\n",
              "      <td>founder</td>\n",
              "      <td>Hg is a leading investor in tax and accounting...</td>\n",
              "      <td>Tim Vandecasteele, co-founder added: \"We want ...</td>\n",
              "      <td>2063</td>\n",
              "      <td>founded_by</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  end_idx                          entities  ... start_idx    string_id\n",
              "0    1024         [Lilium, Baillie Gifford]  ...      1013  invested_in\n",
              "1    1762              [Facebook ’s, Giphy]  ...      1751  acquired_by\n",
              "2    2784    [Global-e, Vitruvian Partners]  ...      2774  invested_in\n",
              "3     680  [Joris Van Der Gucht, Silverfin]  ...       673   founded_by\n",
              "4    2070    [Tim Vandecasteele, Silverfin]  ...      2063   founded_by\n",
              "\n",
              "[5 rows x 8 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KOQ0StVfOEw0"
      },
      "source": [
        "Let's create 2 dictionaries, one that maps each label to a unique integer, and one that does it the other way around."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vDTHyKq2A8-d"
      },
      "source": [
        "id2label = dict()\n",
        "for idx, label in enumerate(df.string_id.value_counts().index):\n",
        "  id2label[idx] = label"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIv9xV_eOthR"
      },
      "source": [
        "As we can see, there are 7 labels (7 unique relationships):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GhnbVGAZBW4B",
        "outputId": "63e8ebeb-9deb-4886-cd72-ae395938daa8"
      },
      "source": [
        "id2label"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{0: 'founded_by',\n",
              " 1: 'acquired_by',\n",
              " 2: 'invested_in',\n",
              " 3: 'CEO_of',\n",
              " 4: 'subsidiary_of',\n",
              " 5: 'partners_with',\n",
              " 6: 'owned_by'}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_ecDeeGVBYDM",
        "outputId": "d9e17f01-d71f-46cc-b2c1-1e0f930be313"
      },
      "source": [
        "label2id = {v:k for k,v in id2label.items()}\n",
        "label2id"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'CEO_of': 3,\n",
              " 'acquired_by': 1,\n",
              " 'founded_by': 0,\n",
              " 'invested_in': 2,\n",
              " 'owned_by': 6,\n",
              " 'partners_with': 5,\n",
              " 'subsidiary_of': 4}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h121Qt799Pc4",
        "outputId": "83ae7d00-3861-416e-d39e-01062feae2c1"
      },
      "source": [
        "df.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(12031, 8)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q9kuc0hAD-ei"
      },
      "source": [
        "## Define the PyTorch dataset and dataloaders\n",
        "\n",
        "Next, we define regular PyTorch datasets and corresponding dataloaders. In PyTorch, you need to define a `Dataset` class that inherits from `torch.utils.data.Dataset`, and you need to implement 3 methods: the `init` method (for initializing the dataset with data), the `len` method (which returns the number of elements in the dataset) and the `getitem()` method, which returns a single item from the dataset.\n",
        "\n",
        "In our case, each item of the dataset consists of a sentence, the spans of 2 entities in the sentence, and a label of the relationship. We use `LukeTokenizer` (available in the Transformers library) to turn these into the inputs expected by the model, which are `input_ids`, `entity_ids`, `attention_mask`, `entity_attention_mask` and `entity_position_ids`.\n",
        "\n",
        "For more information regarding these inputs, refer to the [docs](https://huggingface.co/transformers/model_doc/luke.html#lukeforentitypairclassification) of `LukeForEntityPairClassification`.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jng8JORh65o_"
      },
      "source": [
        "from transformers import LukeTokenizer\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "import torch\n",
        "\n",
        "tokenizer = LukeTokenizer.from_pretrained(\"studio-ousia/luke-base\", task=\"entity_pair_classification\")\n",
        "\n",
        "class RelationExtractionDataset(Dataset):\n",
        "    \"\"\"Relation extraction dataset.\"\"\"\n",
        "\n",
        "    def __init__(self, data):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            data : Pandas dataframe.\n",
        "        \"\"\"\n",
        "        self.data = data\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        item = self.data.iloc[idx]\n",
        "\n",
        "        sentence = item.sentence\n",
        "        entity_spans = [tuple(x) for x in item.entity_spans]\n",
        "\n",
        "        encoding = tokenizer(sentence, entity_spans=entity_spans, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
        "\n",
        "        for k,v in encoding.items():\n",
        "          encoding[k] = encoding[k].squeeze()\n",
        "\n",
        "        encoding[\"label\"] = torch.tensor(label2id[item.string_id])\n",
        "\n",
        "        return encoding"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wWzXmK9BPl6L"
      },
      "source": [
        "Here we instantiate the class defined above with 3 objects: a training dataset, a validation dataset and a test set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9OO_h4LF_uJm"
      },
      "source": [
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "train_df, test_df = train_test_split(df, test_size=0.2, random_state=42, shuffle=True)\n",
        "train_df, val_df = train_test_split(train_df, test_size=0.2, random_state=42, shuffle=False)\n",
        "\n",
        "# define the dataset\n",
        "train_dataset = RelationExtractionDataset(data=train_df)\n",
        "valid_dataset = RelationExtractionDataset(data=val_df)\n",
        "test_dataset = RelationExtractionDataset(data=test_df)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wrdnKYYsBpYW",
        "outputId": "e989ff15-a54e-4107-9236-e68be2e74cfd"
      },
      "source": [
        "train_dataset[0].keys()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['input_ids', 'entity_ids', 'entity_position_ids', 'attention_mask', 'entity_attention_mask', 'label'])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zeutuu_HCQWO"
      },
      "source": [
        "Let's define the corresponding dataloaders (which allow us to iterate over the elements of the dataset):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jL_iys4bCSiQ"
      },
      "source": [
        "train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)\n",
        "valid_dataloader = DataLoader(valid_dataset, batch_size=2)\n",
        "test_dataloader = DataLoader(test_dataset, batch_size=2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lDe3kE85Rec_"
      },
      "source": [
        "Let's verify an example of a batch:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 137
        },
        "id": "LLAQbg92Cea6",
        "outputId": "7a14a53e-2a16-4129-aa7c-a7847055bcba"
      },
      "source": [
        "batch = next(iter(train_dataloader))\n",
        "tokenizer.decode(batch[\"input_ids\"][1])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'<s>The <ent>  Data Security Council of India <ent>, in partnership with <ent2>  PayPal <ent2>, said in an August report that the number of people shopping online in India grew at 73 per cent for big cities and at 400 per cent in smaller towns.\\n\\n</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "mPyZQ9nKRur5",
        "outputId": "0fdf20eb-bc21-443a-e8e6-f7a2d829e2ae"
      },
      "source": [
        "id2label[batch[\"label\"][1].item()]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'partners_with'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ao4iLfbYsktf"
      },
      "source": [
        "## Define a PyTorch LightningModule\n",
        "\n",
        "Let's define the model as a PyTorch LightningModule. A `LightningModule` is actually an `nn.Module`, but with some extra functionality.\n",
        "\n",
        "For more information regarding how to define this, see the [docs](https://pytorch-lightning.readthedocs.io/en/latest/?_ga=2.56317931.1395871250.1622709933-1738348008.1615553774) of PyTorch Lightning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gjmobNfKDOqv"
      },
      "source": [
        "from transformers import LukeForEntityPairClassification, AdamW\n",
        "import pytorch_lightning as pl\n",
        "\n",
        "class LUKE(pl.LightningModule):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.model = LukeForEntityPairClassification.from_pretrained(\"studio-ousia/luke-base\", num_labels=len(label2id))\n",
        "\n",
        "    def forward(self, input_ids, entity_ids, entity_position_ids, attention_mask, entity_attention_mask):     \n",
        "        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, entity_ids=entity_ids, \n",
        "                             entity_attention_mask=entity_attention_mask, entity_position_ids=entity_position_ids)\n",
        "        return outputs\n",
        "    \n",
        "    def common_step(self, batch, batch_idx):\n",
        "        labels = batch['label']\n",
        "        del batch['label']\n",
        "        outputs = self(**batch)\n",
        "        logits = outputs.logits\n",
        "\n",
        "        criterion = torch.nn.CrossEntropyLoss() # multi-class classification\n",
        "        loss = criterion(logits, labels)\n",
        "        predictions = logits.argmax(-1)\n",
        "        correct = (predictions == labels).sum().item()\n",
        "        accuracy = correct/batch['input_ids'].shape[0]\n",
        "\n",
        "        return loss, accuracy\n",
        "      \n",
        "    def training_step(self, batch, batch_idx):\n",
        "        loss, accuracy = self.common_step(batch, batch_idx)     \n",
        "        # logs metrics for each training_step,\n",
        "        # and the average across the epoch\n",
        "        self.log(\"training_loss\", loss)\n",
        "        self.log(\"training_accuracy\", accuracy)\n",
        "\n",
        "        return loss\n",
        "\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "        loss, accuracy = self.common_step(batch, batch_idx)     \n",
        "        self.log(\"validation_loss\", loss, on_epoch=True)\n",
        "        self.log(\"validation_accuracy\", accuracy, on_epoch=True)\n",
        "\n",
        "        return loss\n",
        "\n",
        "    def test_step(self, batch, batch_idx):\n",
        "        loss, accuracy = self.common_step(batch, batch_idx)     \n",
        "\n",
        "        return loss\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        optimizer = AdamW(self.parameters(), lr=5e-5)\n",
        "        return optimizer\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        return train_dataloader\n",
        "\n",
        "    def val_dataloader(self):\n",
        "        return valid_dataloader\n",
        "\n",
        "    def test_dataloader(self):\n",
        "        return test_dataloader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLtoRrv0FtTt"
      },
      "source": [
        "Let's verify a forward pass on a batch:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "s_CMDGDYFkxs",
        "outputId": "f6c93cde-17f6-4878-d324-dff49a5ed00a"
      },
      "source": [
        "batch = next(iter(valid_dataloader))\n",
        "labels = batch[\"label\"]\n",
        "batch.keys()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['input_ids', 'entity_ids', 'entity_position_ids', 'attention_mask', 'entity_attention_mask', 'label'])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9ITWNMSrGUCw",
        "outputId": "0dd9632d-f69d-4ce8-f84b-c04cb5964d98"
      },
      "source": [
        "batch[\"input_ids\"].shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2, 512])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FKfTrvTBFpRp",
        "outputId": "e4feb6e7-9cd5-478b-8fca-96fb093c983c"
      },
      "source": [
        "model = LUKE()\n",
        "del batch[\"label\"]\n",
        "outputs = model(**batch)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at studio-ousia/luke-base were not used when initializing LukeForEntityPairClassification: ['embeddings.position_ids']\n",
            "- This IS expected if you are initializing LukeForEntityPairClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing LukeForEntityPairClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of LukeForEntityPairClassification were not initialized from the model checkpoint at studio-ousia/luke-base and are newly initialized: ['classifier.weight']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5slRAOGTR95T"
      },
      "source": [
        "The initial loss should be around -ln(1/number of classes) = -ln(1/7) = 1.95:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "heWdSlfsSBv8",
        "outputId": "bf4d7f9f-58f5-4c20-98cb-ad272d8a07e0"
      },
      "source": [
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "initial_loss = criterion(outputs.logits, labels)\n",
        "print(\"Initial loss:\", initial_loss)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Initial loss: tensor(1.9616, grad_fn=<NllLossBackward>)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fgEOS69zG_gt"
      },
      "source": [
        "## Train the model\n",
        "\n",
        "Let's train the model. We also use early stopping, to avoid overfitting the training dataset. We also log everything to Weights and Biases, which will give us beautiful charts of the loss and accuracy plotted over time.\n",
        "\n",
        "If you haven't already, you can create an account on the [website](https://wandb.ai/site), then log in in a web browser, and run the cell below: "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3IsneQeIO1ea",
        "outputId": "431825d6-f178-4cbf-fddd-bda068750f7c"
      },
      "source": [
        "import wandb\n",
        "\n",
        "wandb.login()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mnielsrogge\u001b[0m (use `wandb login --relogin` to force relogin)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 437,
          "referenced_widgets": [
            "6c739827bd3c48e68775a3951bebe7c9",
            "46f830096d7449c680d69b6111e73dfe",
            "870102d2c71b473e9f57f62de01636d3",
            "2759d4db7a2b435dba9d6c6b923863b2",
            "08e58a87449945819d0d3fe3a64636a5",
            "78993027eda248ccb5753ae2629e8b63",
            "206e3619d1d8476bbef8a5aa9d1c7b5c",
            "8aedf1e4705a4456ba9dbfb9acefae6a",
            "60e547a3f89b4ff6b21d22f1e4a101f8",
            "c32a8d9b12704d42a2d784c59a2c666f",
            "57371a01d21d474887f025f4a7579acf",
            "f4528f4bd08a41f7966768850ef2e6d0",
            "fcf36c3be3044f38a2ee84c5c204e303",
            "07ce5de710e446649fbd5dc84d393cb0",
            "9e8ed2c19561424c972bac52133019bf",
            "1069cb96cfc7413c86e49569d38e946e",
            "d5d713d04c26461fb098dd15145084c2",
            "69310f37f6f8416cb93426339816405b",
            "80d1309b5e0e4c8e8433402ec3faeb78",
            "62b8e61826804383a726c459852d47f9",
            "390bf45600c2402abaf8baa1459ee4f9",
            "5db797f8ba184f718a11b72b5e798319",
            "3d09aea45dfc43f6ac7f87ed29bf6918",
            "152564b90b62466aab21e6f7744d7024",
            "a22625735f704a8f9751e02b11edd52a",
            "7f3319a4c4bc45479342a111e47d7afb",
            "e879d2939d1048acb218ca36661ae997",
            "6bcc6aa979ce4842a41aa7733a4766c8",
            "c4cf2946d1e3425b8c041e901f9580d9",
            "464a61cb2b4b4eccbeb61fdea07de12c",
            "ad9ba418e52248a09313c8f13ac7bfe1",
            "d7eac30805ce40bb95e4c679ba4126af",
            "466d5d0c87a0493c9a77e4c01233d94e",
            "79d771361fb84c62a18c93f9e70946cf",
            "d8d4525a722a4ba686037016eb8524da",
            "35e227d4e3af46c5856c32587ff18860",
            "12d42fdf1c6e4050ad48907759d3dbea",
            "291b1496ede44d738f63f94c976f6f25",
            "4ce1f43afd33491fa22dd77403cf08ba",
            "0b5233fe3cc047339d33cb2b91ec0908",
            "4efd9113c15d443aa5fadb039e0a31d3",
            "ddadd8532e3a4ca9925d748144f020bd",
            "3883b380d960454ab297f64914d8af25",
            "c3a8c773516c46139dae4dd17e932528",
            "6f83ac1276254e939157238317531e24",
            "b8d8686696334b1089ae1356c17e51c4",
            "eb8531165ba246fe9afb40c498de7d17",
            "b1c31fbb42ea48f3bbe19820f868991e"
          ]
        },
        "id": "ggwuxviZHBNz",
        "outputId": "9212e481-eee2-493b-9d52-8c64ed26f0e6"
      },
      "source": [
        "from pytorch_lightning import Trainer\n",
        "from pytorch_lightning.loggers import WandbLogger\n",
        "from pytorch_lightning.callbacks import EarlyStopping\n",
        "\n",
        "wandb_logger = WandbLogger(name='luke-first-run-12000-articles-bis', project='LUKE')\n",
        "# for early stopping, see https://pytorch-lightning.readthedocs.io/en/1.0.0/early_stopping.html?highlight=early%20stopping\n",
        "early_stop_callback = EarlyStopping(\n",
        "    monitor='val_loss',\n",
        "    patience=2,\n",
        "    strict=False,\n",
        "    verbose=False,\n",
        "    mode='min'\n",
        ")\n",
        "\n",
        "trainer = Trainer(gpus=1, logger=wandb_logger, callbacks=[EarlyStopping(monitor='validation_loss')])\n",
        "trainer.fit(model)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "GPU available: True, used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Tracking run with wandb version 0.10.31<br/>\n",
              "                Syncing run <strong style=\"color:#cdcd00\">luke-first-run-12000-articles-bis</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://wandb.ai/nielsrogge/LUKE\" target=\"_blank\">https://wandb.ai/nielsrogge/LUKE</a><br/>\n",
              "                Run page: <a href=\"https://wandb.ai/nielsrogge/LUKE/runs/19e81nbe\" target=\"_blank\">https://wandb.ai/nielsrogge/LUKE/runs/19e81nbe</a><br/>\n",
              "                Run data is saved locally in <code>/content/wandb/run-20210603_114904-19e81nbe</code><br/><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "  | Name  | Type                            | Params\n",
            "----------------------------------------------------------\n",
            "0 | model | LukeForEntityPairClassification | 274 M \n",
            "----------------------------------------------------------\n",
            "274 M     Trainable params\n",
            "0         Non-trainable params\n",
            "274 M     Total params\n",
            "1,098.052 Total estimated model params size (MB)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "6c739827bd3c48e68775a3951bebe7c9",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "60e547a3f89b4ff6b21d22f1e4a101f8",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "d5d713d04c26461fb098dd15145084c2",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a22625735f704a8f9751e02b11edd52a",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "466d5d0c87a0493c9a77e4c01233d94e",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "4efd9113c15d443aa5fadb039e0a31d3",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/distributed.py:69: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...\n",
            "  warnings.warn(*args, **kwargs)\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U-zsH2M9BUkY"
      },
      "source": [
        "trainer.test()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v5LZj0udEVnh"
      },
      "source": [
        "## Evaluation\n",
        "\n",
        "Instead of calling `trainer.test()`, we can also manually evaluate the model on the entire test set:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "61bca5692a0146ebad95a0b43cc50eaa",
            "23a961a7ebc241b393e408b04d9e5d1d",
            "ff3c708e85b14e5dbbd3cbebc8f1c6cd",
            "e17ac0baf40943c5b0ba3c329c22e992",
            "68fb7646daf141ad9eb00159d7ae8f40",
            "9e9c310829134f8b9825c1a899e88746",
            "ce141efe165d4aa0817c0e7e518039df",
            "42f47d64f2924e7693e16d11a5457136"
          ]
        },
        "id": "MiHWMKvIDW5b",
        "outputId": "78a671ae-7920-41d9-c497-f88c2b36a285"
      },
      "source": [
        "from tqdm.notebook import tqdm\n",
        "from sklearn.metrics import accuracy_score\n",
        "\n",
        "loaded_model.model.eval()\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "loaded_model.to(device)\n",
        "\n",
        "predictions_total = []\n",
        "labels_total = []\n",
        "for batch in tqdm(test_dataloader):\n",
        "    # get the inputs;\n",
        "    labels = batch[\"label\"]\n",
        "    del batch[\"label\"]\n",
        "\n",
        "    # move everything to the GPU\n",
        "    for k,v in batch.items():\n",
        "      batch[k] = batch[k].to(device)\n",
        "\n",
        "    # forward pass\n",
        "    outputs = loaded_model.model(**batch)\n",
        "    logits = outputs.logits\n",
        "    predictions = logits.argmax(-1)\n",
        "    predictions_total.extend(predictions.tolist())\n",
        "    labels_total.extend(labels.tolist())"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "61bca5692a0146ebad95a0b43cc50eaa",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=1204.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3wSCI4zSFFFQ",
        "outputId": "988ce169-44ff-4341-9bf9-efee64dca19c"
      },
      "source": [
        "print(\"Accuracy on test set:\", accuracy_score(labels_total, predictions_total))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Accuracy on test set: 0.9713336103032821\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1hPQJhQLYmen"
      },
      "source": [
        "## Inference\n",
        "\n",
        "Here we test the trained model on a new, unseen sentence."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OThN3tw7A0WA",
        "outputId": "9abf4cb8-fe6f-49ad-9be3-36c43a86773c"
      },
      "source": [
        "loaded_model = LUKE.load_from_checkpoint(checkpoint_path=\"/content/drive/Shareddrives/Datascouts/epoch=3-step=7699.ckpt\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at studio-ousia/luke-base were not used when initializing LukeForEntityPairClassification: ['embeddings.position_ids']\n",
            "- This IS expected if you are initializing LukeForEntityPairClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing LukeForEntityPairClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of LukeForEntityPairClassification were not initialized from the model checkpoint at studio-ousia/luke-base and are newly initialized: ['classifier.weight']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "MgHe48c5DqUf",
        "outputId": "f9af4d47-fe68-4759-b625-c6de26866f63"
      },
      "source": [
        "test_df.iloc[0].sentence"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'Roxanne also co-founded StartHer'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gZNhwFgtYnf_",
        "outputId": "cac9e969-ed3a-4b38-f382-92514e994efd"
      },
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "idx = 2\n",
        "text = test_df.iloc[idx].sentence\n",
        "entity_spans = test_df.iloc[idx].entity_spans  # character-based entity spans \n",
        "entity_spans = [tuple(x) for x in entity_spans]\n",
        "\n",
        "inputs = tokenizer(text, entity_spans=entity_spans, return_tensors=\"pt\")\n",
        "\n",
        "outputs = loaded_model.model(**inputs)\n",
        "logits = outputs.logits\n",
        "predicted_class_idx = logits.argmax(-1).item()\n",
        "print(\"Sentence:\", text)\n",
        "print(\"Ground truth label:\", test_df.iloc[idx].string_id)\n",
        "print(\"Predicted class idx:\", id2label[predicted_class_idx])\n",
        "print(\"Confidence:\", F.softmax(logits, -1).max().item())"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Sentence: One Twitter user suggested that Sun ask Buffett if he regrets not having invested early on in Apple Microsoft Amazon Google Netflix and\n",
            "Ground truth label: invested_in\n",
            "Predicted class idx: invested_in\n",
            "Confidence: 0.9971879124641418\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}